import torch
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
from torch.utils.data.dataset import Dataset
from glob import glob
import cv2
import numpy as np
from torchvision import transforms
from torch.utils.data import DataLoader
from Covariance_Pooling_Steganalytic_Network import Net
import logging
class MyCOCOFolder(Dataset):
    def __init__(self, dir, transform=None):
        self.transform = transform
        self.cover_dir = os.path.join(dir, 'cover')
        self.secret_dir = os.path.join(dir, 'secret')
        self.secret_rev_dir = os.path.join(dir, 'secret_rev')
        self.stego_dir = os.path.join(dir, 'stego')

        # 假设四个目录文件名相同，使用cover_list作为参考
        self.cover_list = [os.path.basename(x) for x in glob(os.path.join(self.cover_dir, '*'))]
        assert len(self.cover_list) != 0, "cover_dir is empty"

    def __getitem__(self, index):
        file_index = int(index)

        return self.cover_list[file_index]


    def __len__(self):
        return len(self.cover_list)
    
def main():
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = False

    batch_size = 128
    DATA_DIR = '/home/zzc/hack/dataset'

    model =Net().cuda()
    state_dict = torch.load('/home/zzc/hack/00010_dec.pth')
    model.load_state_dict(state_dict['model_state_dict'])
    # state_dict = torch.load('/home/zzc/hack/Cov-Pooling-Steganalytic-Network/cover_stego_save_models/e_J_d_covnet.pt')
    # model.load_state_dict(state_dict['original_state'])
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
        ]
    )

    model_data_path = os.path.join(DATA_DIR, 'A_ete', 'train') 
    dataset = MyCOCOFolder(dir=model_data_path, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, drop_last=True)


    accuracy_list = []
  # self.eval()

    model_name_list_meta_test = ['A_ete', 'J_udh','D_cci', 'E_unet', 'F_dca',
                          'H_hinet', 'I_resnet', 'G_ete_dcva', 'B_hps', 'C_stegnet']

    indexs = []

    model_data_path = os.path.join(DATA_DIR, 'A_ete', 'val') 
    dataset = MyCOCOFolder(dir=model_data_path, transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    for images_indexs in dataloader:
      indexs = images_indexs
      break
    for i, model_name in enumerate(model_name_list_meta_test):
        correct = 0
        model_data_path = os.path.join(DATA_DIR, model_name, 'val')
        secret_images = []
        rev_images = []

        for index in indexs:

            secret_path = os.path.join(model_data_path, 'secret', index)
            rev_path = os.path.join(model_data_path, 'rev_s', index)


            # 读取图像
            secret_data = torch.from_numpy(np.load(secret_path)).cuda()
            rev_data = torch.from_numpy(np.load(rev_path)).cuda()

            secret_images.append(secret_data.cuda())
            rev_images.append(rev_data.cuda())


        # 将 stego_images 列表转换为张量
        secret_images = torch.stack(secret_images)
        rev_images = torch.stack(rev_images)

        secret_labels = torch.ones(secret_images.size(0), dtype=torch.long).cuda()
        rev_labels = torch.zeros(rev_images.size(0), dtype=torch.long).cuda()

        data = torch.stack([secret_images, rev_images])
        data = data.reshape(data.shape[0]*data.shape[1], data.shape[2], data.shape[3], data.shape[4])
        label = torch.stack([secret_labels, rev_labels])
        label = label.reshape(-1)
        output = model(data)

        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(label.view_as(pred)).sum().item()
        accuracy = correct / (len(secret_images) * 2)

        # Log individual model performance
        print(f"Accuracy for model {model_name}: {accuracy:.6f}")


        accuracy_list.append(accuracy)


    accuracy_mean = sum(accuracy_list) / len(accuracy_list)
  

    print("____________________________________________")
    print(f"Avearage Accuracy for all models: {accuracy_mean:.6f}")
    print("____________________________________________")


if __name__ == '__main__':
    main()